import torch
import matplotlib.pyplot as plt

# p = torch.randn((R, N))
# print("p.shape", p.shape)
# print("a.shape", a.shape)
# print("torch.mm(p,a).shape", torch.mm(p, a).shape)
"""
欲将一组N维向量降到K维（0<K<N），其目标是选择K个单位正交基，使得原始数据变换到这组基上，各维度间协方差为0，而各维度自身的方差尽可能大，并且最终只取最大的K个方差。
设原始数据集有M条N维数据，欲将其降维到K维，则PCA降维算法步骤如下：
（1）将原始数据按列组成N×M的矩阵X；
（2）去中心化;
（3）求矩阵X的协方差矩阵C；
（4）求矩阵C的特征值及对应的特征向量；
（5）取其中K个最大的特征值并从大到小按排列，并将其对应的特征向量按行排列成矩阵P；
（6）Y=PX即为降维后的数据集（按列排列了，矩阵大小为K×M）。
"""


class PCA_tools(object):
    def __init__(self, input, k=2):
        self.input = input
        self.k = k
        self.PCA_value = self.PCA()

    def PCA(self):
        data_new = self.input - self.input.mean(axis=0)  # 去中心化
        # print("去中心化", data_new.shape)
        # H1 = data_new.T @ data_new / (data_new.shape[0] - 1)  # 计算协方差矩阵，矩阵相乘后除以[点数-1]
        # print(H1)
        H = torch.cov(self.input.T)  # 计算结果维度和输入数据行数保持一致，因此输入数据要进行转置（np.cov输入数据去不去中心化结果一样）
        # print("求矩阵X的协方差矩阵H", H)
        # # 使用ED分解
        # eigen_vals, eigen_vecs = np.linalg.eig(H)
        # sort_index = np.argsort(eigen_vals)[::-1]
        # select_eigen_vecs = eigen_vecs[:, sort_index[:k]]
        # 使用SVD分解
        U, S, V = torch.linalg.svd(H)
        # print("U", U)
        # print("S", S)
        # print("V", V)
        """
        H = U@S组成对角阵@T
        """
        values, sort_index = torch.sort(S, descending=True)
        select_eigen_vecs = U[:, sort_index[:self.k]]  # 求前k大的奇异值对应的特征向量
        # 计算PCA后的结果，并可视化
        points_after_PCA = torch.mm(data_new, select_eigen_vecs)  # 要使用去中心化的数据进行运算！！！
        return points_after_PCA

    def show_pca(self):
        if self.k == 2:
            self.show_pca_2d()
        else:
            self.show_pca_3d()

    def show_pca_2d(self):
        # 创建一个新的图形
        fig = plt.figure()
        x = self.PCA_value.T[0]
        y = self.PCA_value.T[1]
        # 添加一个3D子图
        ax = fig.add_subplot(111)
        # 绘制3D散点图
        ax.scatter(x, y, c='r', marker='o')
        # 设置标题和轴标签
        ax.set_title('3D散点图示例')
        ax.set_xlabel('X轴')
        ax.set_ylabel('Y轴')
        # 显示图形
        plt.show()

    def show_pca_3d(self):
        # 创建一个新的图形
        fig = plt.figure()
        x = self.PCA_value.T[0]
        y = self.PCA_value.T[1]
        z = self.PCA_value.T[2]
        # 添加一个3D子图
        ax = fig.add_subplot(111, projection='3d')
        # 绘制3D散点图
        ax.scatter(x, y, z, c='r', marker='o')
        # 设置标题和轴标签
        ax.set_title('3D散点图示例')
        ax.set_xlabel('X轴')
        ax.set_ylabel('Y轴')
        ax.set_zlabel('Z轴')
        # 显示图形
        plt.show()


"""
M = 400
N = 32
# PCA模块处理方法
input = torch.randn(N, M)
K = 2  # PCA降维目标维度
pca = PCA_tools(input, K)
pca.show_pca()
"""
